#' Generate Adjacency Matrix
#'
#' @description Generate an adjacency matrix representing the containment
#' relationships between clusters using the inferred evolutionary tree
#' obtained from the post-clustering processing steps.
#'
#' @param inferTreeInput A data.frame, which is the 'inferTree' data.frame from
#' the list returned by the postProcess() function, containing the mutation
#' information and CCF for each sample.
#'
#' @return Returns a list containing a data.frame for the adjacency matrix and
#' a data.frame for the CCF of clusters.
#'
#' @export
GenerateAdjMatrix <- function(inferTreeInput) {
  if (length(unique(inferTreeInput$cluster)) < 2) {
    cat("\nOnly one cluster, cannot infer tree\n")
    return(NULL)
  }
  # Initialize the adjacency matrix
  adj_matrix <-
    matrix(0,
      nrow = length(unique(inferTreeInput$cluster)),
      ncol = length(unique(inferTreeInput$cluster))
    )

  # Generate the adjacency matrix sample by sample
  for (i in 1:(ncol(inferTreeInput) - 4)) {
    for (j in 1:length(unique(inferTreeInput$cluster))) {
      for (k in 1:length(unique(inferTreeInput$cluster))) {
        if (j != k) {
          cluster_j <-
            unlist(inferTreeInput[inferTreeInput$cluster == j, -c(
              1, 2, 3,
              ncol(inferTreeInput)
            )][i])
          cluster_k <-
            unlist(inferTreeInput[inferTreeInput$cluster == k, -c(
              1, 2, 3,
              ncol(inferTreeInput)
            )][i])
          if (mean(cluster_j) >= mean(cluster_k)) {
            adj_matrix[j, k] <- adj_matrix[j, k] + 1
          }
        }
      }
    }
  }

  # Dealing with points that violate the crossover rule
  for (j in 1:length(unique(inferTreeInput$cluster))) {
    for (k in 1:length(unique(inferTreeInput$cluster))) {
      if (adj_matrix[j, k] != 0 && adj_matrix[k, j] != 0) {
        adj_matrix[j, k] <- 0
        adj_matrix[k, j] <- 0
      }
    }
  }

  # Obtain the average ccf for each cluster
  cluster_ccf <- data.frame(cluster = numeric(), ccf = numeric())
  for (i in 1:length(unique(inferTreeInput$cluster))) {
    cluster_ccf <-
      rbind(
        cluster_ccf,
        data.frame(
          cluster = i,
          ccf = mean(unlist(
            inferTreeInput[inferTreeInput$cluster == i, -c(1, 2, 3, ncol(inferTreeInput))]
          ))
        )
      )
  }

  cat("\nAdjacency Matrix has been Generated\n")
  return(list(adj_matrix = adj_matrix, cluster_ccf = cluster_ccf))
}


#' Finding the Maximum Sum Sub Node
#'
#' @description Find the child node that does not exceed the maximized ccf sum
#' of the parent node.
#'
#' @param adj_matrix The adjacency matrix generated by GenerateAdjMatrix().
#' @param parent_index The ID of the parent node cluster.
#' @param error_buf Size of the error buffer.
#'
#' @return Returns a vector of child node IDs, where the child nodes are
#' children of the largest sum of ccf's of that parent node.
findMaxSumChild <-
  function(adj_matrix, parent_index, error_buf) {
    parent_ccf <-
      adj_matrix$cluster_ccf[adj_matrix$cluster_ccf$cluster == parent_index, ]$ccf

    child_index <- which(adj_matrix$adj_matrix[parent_index, ] != 0)

    child_ccf <-
      adj_matrix$cluster_ccf[child_index, ]$ccf

    names(child_ccf) <- child_index

    sorted_child_ccf <- sort(child_ccf, decreasing = TRUE)

    children <- vector()

    child_ccf_sum <- 0

    for (i in seq_along(sorted_child_ccf)) {
      child_ccf_sum <- child_ccf_sum + sorted_child_ccf[i]
      if (child_ccf_sum <= (parent_ccf + parent_ccf * error_buf)) {
        children <- c(children, as.numeric(names(sorted_child_ccf)[i]))
      } else {
        child_ccf_sum <- child_ccf_sum - sorted_child_ccf[i]
      }
    }

    return(children)
  }



#' Generating the Default Tree
#'
#' @description Quickly find a clonal evolutionary tree from the adjacency
#' matrix that meets the sum conditions and crossing rules based on a greedy
#' strategy.
#'
#' @param adj_matrix A data.frame of adjacency matrix generated by
#' GenerateAdjMatrix(), representing the containment relationship between clusters.
#' @param error_buf Size of the error buffer, default is 0.1.Represent the CCF
#' of the parent cluster as the CCF of the parent cluster multiplied by
#' 'error_buf' to account for error noise.
#'
#' @return Returns a data.frame where the first column represents the parent
#' cluster and the second column represents the child clusters.
#'
#' @export
#'
#' @importFrom igraph graph_from_adjacency_matrix get.edgelist
generateDefaultTree <- function(adj_matrix, error_buf = 0.1) {
  # Initialize the default tree
  n <- nrow(adj_matrix$adj_matrix)
  default_tree <- matrix(
    data = 0,
    nrow = n,
    ncol = n
  )

  # Create a queue to store pending nodes
  queue <- 1

  while (length(queue) > 0) {
    current_index <- queue[1]
    queue <- queue[-1]

    children <-
      findMaxSumChild(adj_matrix, current_index, error_buf)

    for (child in children) {
      if (sum(default_tree[, child]) == 0) {
        default_tree[current_index, child] <- 1
        adj_matrix$adj_matrix[, child] <- 0
      }

      queue <- c(queue, child)
    }
  }

  # Convert adjacency matrix to directed graph object
  graph <-
    graph_from_adjacency_matrix(default_tree, mode = "directed")

  edges <- get.edgelist(graph)

  default_tree <- as.data.frame(edges, stringsAsFactors = FALSE)

  colnames(default_tree) <- c("parent", "child")

  cat("\nDefault Tree has been Generated\n")

  return(default_tree)
}


#' Check That The Sum Condition is Satisfied
#'
#' @param edge_df Cloned evolutionary tree data.frame with parent node in the
#' first column and child nodes in the second column.
#' @param cluster_ccf CCF for each cluster.
#' @param error_buf Size of the error buffer.
#'
#' @return Returns a boolean value, if TRUE means the sum condition is satisfied.
isSumRule <- function(edge_df, cluster_ccf, error_buf) {
  for (i in unique(edge_df$parent)) {
    child_index <- edge_df[edge_df$parent == i, ]$child
    child_ccf <- cluster_ccf[child_index, ]$ccf
    parent_ccf <- cluster_ccf[i, ]$ccf
    if (sum(child_ccf) > (parent_ccf + parent_ccf * error_buf)) {
      return(FALSE)
    }
  }
  return(TRUE)
}


#' Enumerate and Sort All Spanning Trees
#'
#' @description Based on the adjacency matrix, enumerate all the cloned
#' evolutionary trees that satisfy the sum condition and the crossover rule,
#' and sort them according to the probability of occurrence of the total edges.
#'
#' @param adj_matrix A data.frame of adjacency matrix generated by
#' GenerateAdjMatrix(),representing the containment relationship between clusters.
#' @param error_buf Size of the error buffer,default is 0.1. Represent the CCF
#' of the parent cluster as the CCF of the parent cluster multiplied by
#' 'error_buf' to account for error noise.
#' @param max_tree_num The maximum number of cloned evolutionary trees to
#' enumerate if there are multiple different evolutionary trees, default is 20.
#' @param remove_cluster A Boolean value indicating whether it is permissible to
#' delete the cluster with the lowest number of SNVs and then re-run the
#' enumeration tree process when the current cluster node is unable to generate
#' an evolutionary tree.
#' @param inferTreeInput A data.frame of 'inferTree' that is consistent with
#' the input data of GenerateAdjMatrix(). Valid only if 'remove_cluster' is TRUE.
#'
#' @return Returns a list that includes a sorted list of evolutionary trees and
#' a sorted weights data.frame, where each data.frame in the evolutionary tree
#' list represents an evolutionary tree.
#'
#' @importFrom dplyr arrange
#'
#' @export
enumerateAllTree <-
  function(adj_matrix,
           error_buf = 0.1,
           max_tree_num = 20,
           remove_cluster = FALSE,
           inferTreeInput = NULL) {
    tree_list <- list()

    generated_trees <- list()

    # The number of snv in the cluster
    if (remove_cluster == T) {
      if (is.null(inferTreeInput)) {
        cat("\nParameter 'inferTreeInput' must not be empty when 'remove_cluster==TRUE'\n")
        return(NULL)
      } else {
        snv_num <- table(inferTreeInput$cluster)
        snv_num <- sort(snv_num)
      }
    }

    while (TRUE) {
      m <- adj_matrix$adj_matrix
      m[m != 0] <- 1
      m <- as.integer(prod(colSums(m)[-1]))
      if (m > 10000) {
        m <- 10000
      }
      # When all spanning trees have been enumerated or the trees that satisfy
      # the sum rule exceed max_tree_num, the loop ends.
      while (length(generated_trees) < m &&
        length(tree_list) < max_tree_num) {
        edge_df <- data.frame(parent = numeric(), child = numeric())

        for (i in 2:ncol(adj_matrix$adj_matrix)) {
          parent <- which(adj_matrix$adj_matrix[, i] != 0)
          parent <- sample(parent, 1)
          edge_df[i - 1, 1] <- parent
          edge_df[i - 1, 2] <- i
        }

        if (!toString(edge_df) %in% generated_trees) {
          generated_trees <- append(generated_trees, toString(edge_df))
        } else {
          next
        }

        if (isSumRule(edge_df, adj_matrix$cluster_ccf, error_buf)) {
          tree_list <- append(tree_list, list(edge_df))
        }
      }

      if (length(tree_list) == 0 && length(snv_num) != 2) {
        c <- as.numeric(names(snv_num[1]))
        adj_matrix$adj_matrix <- adj_matrix$adj_matrix[-c, ]
        adj_matrix$adj_matrix <- adj_matrix$adj_matrix[, -c]
        adj_matrix$cluster_ccf <- adj_matrix$cluster_ccf[-c, ]
        adj_matrix$cluster_ccf <-
          reClusterID(df = adj_matrix$cluster_ccf)
        snv_num <- snv_num[-1]
        generated_trees <- list()
        cat("\ncluster", c, "is removed\n")
      } else {
        break
      }
    }

    if (length(tree_list) == 0) {
      cat("\nFailed To Enumerate Trees, Unable To Find a Matching Spanning Tree\n")
      return(NULL)
    }

    ## Sorting the tree list
    weight_df <-
      data.frame(
        tree = seq_along(tree_list),
        sample_weight = rep(0, length(tree_list)),
        edge_prob_weight = rep(0, length(tree_list))
      )

    for (i in seq_along(tree_list)) {
      weight <- 0
      for (j in 1:(nrow(tree_list[[1]]))) {
        pc <- tree_list[[i]]$parent[j]
        cc <- tree_list[[i]]$child[j]
        weight <- weight + adj_matrix$adj_matrix[pc, cc]
      }
      weight_df$sample_weight[i] <- weight
    }

    for (i in seq_along(tree_list)) {
      sum_edge_prob <- 0
      for (j in 1:(nrow(tree_list[[1]]))) {
        pc <- tree_list[[i]]$parent[j]
        cc <- tree_list[[i]]$child[j]
        count <- 0
        for (k in seq_along(tree_list)) {
          if (any(tree_list[[k]][tree_list[[k]]$parent == pc, ]$child == cc)) {
            count <- count + 1
          }
        }
        prob <- count / length(tree_list)
        sum_edge_prob <- sum_edge_prob + prob
      }
      weight_df$edge_prob_weight[i] <- sum_edge_prob
    }

    weight_df <-
      arrange(weight_df, desc(sample_weight), desc(edge_prob_weight))
    tree_list <- tree_list[weight_df$tree]
    names(tree_list) <- paste0("Tree", seq_along(tree_list))
    weight_df$tree <- paste0("Tree", seq_along(tree_list))

    cat("\nAll Spanning Trees are Enumerated and Sorted\n")

    return(list(tree_list = tree_list, weight_df = weight_df))
  }



#' Calculation of Subclone Ratio
#'
#' @description Calculate the subclone proportion of each sample based on any
#' evolutionary tree.
#'
#' @param clusterMat A data.frame of the subclone clustering results is
#' consistent with the 'mat' in the list returned by postProcess().
#' @param evolution An evolutionary tree data.frame returned by
#' generateDefaultTree() or enumerateAllTree().
#'
#' @return Returns a data.frame in which the proportion of subclones in each
#' cluster for each sample has been calculated.
#'
#' @importFrom tidyr pivot_longer
#' @importFrom dplyr mutate group_by summarise left_join filter
#' @importFrom tidyselect ends_with
#'
#' @export
calcCloneProp <- function(clusterMat, evolution) {
  clusterMat <-
    clusterMat[, names(clusterMat)[c(grep("*.ccf", names(clusterMat)), ncol(clusterMat))]]

  clusterMat <-
    clusterMat[clusterMat$cluster %in% unique(unlist(evolution)), ]

  long_data <- clusterMat %>%
    pivot_longer(
      cols = ends_with(".ccf"),
      names_to = "sample",
      values_to = "ccf"
    ) %>%
    mutate(sample = gsub("\\.ccf", "", sample))

  mean_ccf <- long_data %>%
    group_by(sample, cluster) %>%
    summarise(mean_ccf = mean(ccf), .groups = "drop")

  # Calculate the average ccf of the clusters in each sample relative to the
  # percentage of cluster 1
  mean_ccf <- mean_ccf %>%
    group_by(sample) %>%
    mutate(relative_ccf = mean_ccf / mean_ccf[cluster == 1] * 100)

  # Plus evolutionary relationships between clusters
  mean_ccf <- mean_ccf %>%
    left_join(evolution, by = c("cluster" = "child")) %>%
    mutate(parent = ifelse(is.na(parent), 0, parent))

  # Calculate the sum of the CCFs of the child nodes of each cluster
  # within each sample
  mean_ccf_p <- mean_ccf %>%
    filter(parent != 0) %>%
    group_by(sample, parent) %>%
    summarise(mean_ccf_p = sum(mean_ccf), .groups = "drop")

  # Merge the sum of the CCFs of the child nodes into the original data frame
  mean_ccf <- mean_ccf %>%
    left_join(mean_ccf_p, by = c("sample", "cluster" = "parent")) %>%
    mutate(mean_ccf_p = ifelse(is.na(mean_ccf_p),
      mean_ccf,
      mean_ccf - mean_ccf_p
    ))

  mean_ccf[mean_ccf < 0] <- 0

  cat("\nCalculate Clone Proportion Finished\n")
  return(as.data.frame(mean_ccf))
}


#' Inferring Sample Evolutionary Trees
#'
#' @description Infer the common ancestor and evolutionary relationship between
#' samples based on the grouping returned by runSampleGroup().
#'
#' @param sampleGroup A list of grouping cases returned by runSampleGroup().
#'
#' @return Returns an igraph tree object.
#'
#' @importFrom igraph graph_from_edgelist V V<- degree is_connected
#'
#' @export
generateSampleTree <- function(sampleGroup) {
  labels <- unique(sampleGroup$group$label)

  label_df <- data.frame(
    label = labels,
    ones_count = sapply(labels, function(x) {
      sum(strsplit(x, "")[[1]] == "1")
    })
  )

  label_df <- label_df[order(label_df$ones_count), ]

  edges <-
    data.frame(
      from = character(),
      to = character(),
      stringsAsFactors = FALSE
    )

  for (i in 1:(nrow(label_df) - 1)) {
    to <- label_df$label[i]
    to_ones <- strsplit(to, "")[[1]] == "1"

    for (j in (i + 1):nrow(label_df)) {
      from <- label_df$label[j]
      from_ones <- strsplit(from, "")[[1]] == "1"

      if (all(from_ones[to_ones])) {
        edges <-
          rbind(
            edges,
            data.frame(
              from = from,
              to = to,
              stringsAsFactors = FALSE
            )
          )
      }
    }
  }

  edges_to_delete <-
    data.frame(
      from = character(),
      to = character(),
      stringsAsFactors = FALSE
    )

  for (child in unique(edges$to)) {
    parent_levels <-
      label_df$ones_count[label_df$label %in% edges$from[edges$to == child]]
    min_level <- min(parent_levels)

    for (parent in edges$from[edges$to == child]) {
      if (label_df$ones_count[label_df$label == parent] != min_level) {
        edges_to_delete <-
          rbind(
            edges_to_delete,
            data.frame(
              from = parent,
              to = child,
              stringsAsFactors = FALSE
            )
          )
      }
    }
  }

  edges <-
    edges[!with(
      edges,
      paste(from, to) %in% paste(edges_to_delete$from, edges_to_delete$to)
    ), ]

  g <- graph_from_edgelist(as.matrix(edges), directed = TRUE)

  # Setting the size of the node changes according to the number of mutations
  snv_num <- table(sampleGroup$group$label)
  snv_num <- snv_num[V(g)$name]
  V(g)$label <- snv_num
  snv_num[snv_num > 60] <- 60
  snv_num[snv_num < 30] <- 30
  V(g)$size <- snv_num
  V(g)$label <- paste0(V(g)$name, "\n(", V(g)$label, ")")

  # Add sample names for leaf nodes
  sampleName <-
    names(sampleGroup$group[, -c(1:3, ncol(sampleGroup$group))])
  sampleName <- gsub(".ccf", "", sampleName)
  leaf <- V(g)$name[V(g)[degree(g, mode = "out") == 0]]
  leaf <- grep("^0*1{1}0*$", leaf, value = TRUE)
  for (i in leaf) {
    V(g)$label[V(g)$name == i] <-
      paste0("\n\n", V(g)$label[V(g)$name == i], "\n\n", sampleName[as.numeric(gregexpr("1", i)[[1]])])
  }

  if (is_connected(g)) {
    cat("\nSample Tree has been Generated\n")
  } else {
    cat("\nFailed to Infer sample tree, the tree is not connected\n")
  }

  return(g)
}



#' Save Adjacency Matrix
#'
#' @description Save the adjacency matrix along with the average CCF for each
#' cluster.
#'
#' @param adjmat A list returned by GenerateAdjMatrix().
#' @param saveDir Save directory, automatically save all files to this directory,
#' and automatically create it if it does not exist.
#'
#' @importFrom utils write.table
#'
#' @export
saveAdjmat <- function(adjmat, saveDir) {
  if (!dir.exists(saveDir)) {
    dir.create(saveDir, recursive = T)
  }

  mat <- as.data.frame(adjmat$adj_matrix)
  rownames(mat) <-
    paste0("cluster", 1:nrow(mat))
  colnames(mat) <-
    paste0("cluster", 1:ncol(mat))

  write.table(
    data.frame(ID = rownames(mat), mat),
    paste0(saveDir, "/adj_matrix.tsv"),
    sep = "\t",
    row.names = F
  )

  write.table(
    adjmat$cluster_ccf,
    paste0(saveDir, "/cluster_mean_ccf.tsv"),
    sep = "\t",
    row.names = F
  )
}



#' Save Default Tree
#'
#' @param defaultTree A default tree data.frame returned by generateDefaultTree().
#' @param saveDir Save directory, automatically save all files to this directory,
#' and automatically create it if it does not exist.
#'
#' @importFrom utils write.table
#'
#' @export
saveDefaultTree <- function(defaultTree, saveDir) {
  if (!dir.exists(saveDir)) {
    dir.create(saveDir, recursive = T)
  }

  write.table(
    defaultTree,
    paste0(saveDir, "/default_tree.tsv"),
    sep = "\t",
    row.names = F
  )
}



#' Save All Spanning Trees
#'
#' @description Save all spanning tree and sort weight data frames.
#'
#' @param sortedAllTree A list returned by enumerateAllTree().
#' @param saveDir Save directory, automatically save all files to this directory,
#' and automatically create it if it does not exist.
#'
#' @importFrom utils write.table
#'
#' @export
saveAllTree <- function(sortedAllTree, saveDir) {
  if (!dir.exists(saveDir)) {
    dir.create(saveDir, recursive = T)
  }

  write.table(
    sortedAllTree$tree_list,
    paste0(saveDir, "/all_tree.tsv"),
    sep = "\t",
    row.names = F
  )

  write.table(
    sortedAllTree$weight_df,
    paste0(saveDir, "/all_tree_weight.tsv"),
    sep = "\t",
    row.names = F
  )
}


#' Save Subclone Ratio
#'
#' @param prop A data.frame returned by calcCloneProp().
#' @param saveDir Save directory, automatically save all files to this directory,
#' and automatically create it if it does not exist.
#'
#' @importFrom utils write.table
#'
#' @export
saveCloneProp <- function(prop, saveDir) {
  if (!dir.exists(saveDir)) {
    dir.create(saveDir, recursive = T)
  }

  write.table(
    prop,
    paste0(saveDir, "/clone_proportion.tsv"),
    sep = "\t",
    row.names = F
  )
}
